package com.manli.api.util;

import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.apache.commons.math3.util.Pair;

import java.util.List;
import java.util.Map;
import java.util.SortedMap;
import java.util.TreeMap;

public class WeightRandom<K,V extends Number> {
    private TreeMap<Double, K> weightMap = new TreeMap<Double, K>();

    public WeightRandom(List<Pair<K, V>> list) {
        Preconditions.checkNotNull(list, "list can NOT be null!");
        for (Pair<K, V> pair : list) {
            Preconditions.checkArgument(pair.getValue().doubleValue() > 0, String.format("非法权重值：pair=%s", pair));

            double lastWeight = this.weightMap.size() == 0 ? 0 : this.weightMap.lastKey().doubleValue();//统一转为double
            this.weightMap.put(pair.getValue().doubleValue() + lastWeight, pair.getKey());//权重累加
        }
    }

    public K random() {
        double randomWeight = this.weightMap.lastKey() * Math.random();
        SortedMap<Double, K> tailMap = this.weightMap.tailMap(randomWeight, false);
        return this.weightMap.get(tailMap.firstKey());
    }

    /**
     * 例子，new Pair(”需要返回的数值“, ”该数字的权重“)
     * @param args
     */
    public static void main(String[] args) {
        List<Pair<String, Integer>> list = Lists.newArrayList();
        list.add(new Pair("A", 5));
        list.add(new Pair("B", 1));
        list.add(new Pair("C", 3));
        list.add(new Pair("D", 4));
//        list.add(new Pair("E", 0));
        WeightRandom<String, Integer> random = new WeightRandom(list);

        Map<String, Integer> countMap = Maps.newHashMap();
        for (int i = 0; i < 100000000; i++) {

            String randomKey = random.random();
            countMap.put(randomKey, countMap.getOrDefault(randomKey, 0) + 1);
        }

        for (Pair<String, Integer> pair : list) {
            System.out.println("{" + pair.getKey()+ "}:{" + countMap.get(pair.getKey()) + "}");
        }
    }

}
